#include "container.h"
#include "container_config.h"
#include "file_helper.hpp"
#include "fmt/core.h"
#include <algorithm>
#include <cerrno>
#include <csignal>
#include <cstddef>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fmt/color.h>
#include <fmt/core.h>
#include <fstream>
#include <new>
#include <sched.h>
#include <string>
#include <string_view>
#include <sys/mount.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <system_error>
#include <unistd.h>

constexpr std::string_view kCgroupFolder = "/sys/fs/cgroup";
constexpr std::string_view kCgroupCpuFolder = "/sys/fs/cgroup/cpu/container";
constexpr std::string_view kCgroupPidFolder = "/sys/fs/cgroup/pids/container";
constexpr std::string_view kCgroupMemoryFloder =
    "/sys/fs/cgroup/memory/container";

void write_rule(std::string_view file_path, std::string_view value);
void write_rule(std::string_view file_path, std::integral auto value) {
  write_rule(file_path, std::to_string(value));
}

std::string getExecutablePath() {
  char result[PATH_MAX];
  ssize_t count = readlink("/proc/self/exe", result, PATH_MAX);
  if (count == -1) {
    panic("Failed to get executable path!");
  }
  return std::string(result, count);
}

std::string getParentPath(std::string_view path) {
  std::filesystem::path file_path{path};
  return file_path.parent_path().string();
}

void unzipImage(std::string_view work_dir, std::string_view image_arch_name) {
  printInfoMessage(fmt::format("Decompression in progress, current image: {}",
                               image_arch_name));

  std::filesystem::path images_dir(fmt::format("{}/images", work_dir));

  if (!std::filesystem::exists(images_dir)) {
    std::error_code ec;
    auto succ = std::filesystem::create_directory(images_dir, ec);
    if (!succ) {
      panic(ec.message());
    }
  }

  int ret = 0;
  if (image_arch_name.ends_with(".zip")) {
    if (programExists("unzip")) {
      ret = system(fmt::format("unzip {}", image_arch_name).data());
    } else {
      panic("Unzip does not exist!");
    }
  } else if (image_arch_name.ends_with(".tar.gz")) {
    if (programExists("tar")) {
      ret = system(
          fmt::format("tar -zxf {} -C {}/images", image_arch_name, work_dir)
              .data());
    } else {
      panic("Tar does not exist!");
    }
  } else {
    panic("Unknown file type! image file should be: zip/gz");
  }

  if (ret != 0) {
    panic("Failed to decompress image file!");
  }

  printInfoMessage("Decompression success!");
}

char *alloc_memory(size_t stack_size) {
  auto *stack = new (std::nothrow) char[stack_size];
  if (stack == nullptr) {
    panic("Insufficient memory to allocate more space!");
  }

  return stack;
}

void write_rule(std::string_view file_path, std::string_view value) {
  std::fstream s(file_path.data(), std::fstream::out);

  if (!s.is_open()) {
    panic(fmt::format("open rule file faild! reason: {}", strerror(errno)));
  }

  s.write(value.data(), value.length());

  s.close();
}

void Container::limitProcessCreation() {
  printInfoMessage("Environmental preparation in progress");
  mkdir(kCgroupPidFolder.data(), S_IRUSR | S_IWUSR);

  const std::string pid = std::to_string(getpid());

  write_rule(fmt::format("{}/cgroup.procs", kCgroupPidFolder), pid);
  write_rule(fmt::format("{}/notify_on_release", kCgroupPidFolder), "1");
  write_rule(fmt::format("{}/pids.max", kCgroupPidFolder), this->conf.pids_limit);

  mkdir(kCgroupCpuFolder.data(), S_IRUSR | S_IWUSR);
  write_rule(fmt::format("{}/cgroup.procs", kCgroupCpuFolder), pid);
  write_rule(fmt::format("{}/cpu.cfs_period_us", kCgroupCpuFolder), "100000");
  write_rule(fmt::format("{}/cpu.cfs_quota_us", kCgroupCpuFolder),
             this->conf.cpu_limit);

  mkdir(kCgroupMemoryFloder.data(), S_IRUSR | S_IWUSR);
  write_rule(fmt::format("{}/tasks", kCgroupMemoryFloder), pid);
  write_rule(fmt::format("{}/memory.limit_in_bytes", kCgroupMemoryFloder),
             this->conf.memory_limit);

  printInfoMessage("Environmental preparation completed");
}

template <typename... P> int run(P... params) {
  char *args[] = {(char *)params..., NULL};
  return execvp(args[0], args);
}

void setEnv() {
  clearenv();
  setenv("TERM", "xterm-256color", 0);
  setenv("PATH", "/bin/:/sbin/:usr/bin:/usr/sbin", 0);
}

void setDir(Container* container_ptr) {
  chroot(fmt::format("{}/images", container_ptr->getWorkDir()).data());
  chdir("/");
}

int container_run(void *arg) {
  auto conatiner_ptr = static_cast<Container *>(arg);

  printInfoMessage(
      fmt::format("container is running! child process: {}", getpid()));

  setEnv();
  setDir(conatiner_ptr);

  mount("proc", "/proc", "proc", 0, 0);

  auto runSubprocess = +[](void *args) -> int {
    run("/bin/sh");
    return 0;
  };
  clone(runSubprocess, alloc_memory(65535) + 65535, SIGCHLD, 0);

  wait(nullptr);
  umount("/proc");

  return EXIT_SUCCESS;
}

Container::Container(ContainerConfig config) : conf(std::move(config)) {
  auto exec_path = getExecutablePath();
  this->workDir = getParentPath(exec_path);

  auto image_name_begin =
      (std::find(conf.base_image.rbegin(), conf.base_image.rend(), '/').base());
  auto image_name_end = std::find(image_name_begin, conf.base_image.end(), '.');
  this->imageName = std::string{image_name_begin, image_name_end};
  printInfoMessage(fmt::format("image name: {}", imageName));
}

void Container::start() {

  unzipImage(this->workDir, this->conf.base_image);

  limitProcessCreation();
  int ret =
      clone(container_run, alloc_memory(this->kStackSize) + kStackSize,
            SIGCHLD | CLONE_NEWUTS | CLONE_NEWPID, static_cast<void *>(this));
  if (ret == -1) {
    panic(fmt::format("clone child process faild: {}", strerror(errno)));
  }
  wait(nullptr);
}
